package ssh

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"sync"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"
)

// SSHClient SSH客户端接口
type SSHClient interface {
	// Connect 连接到SSH服务器
	Connect(ctx context.Context) error
	// Disconnect 断开连接
	Disconnect() error
	// ExecuteCommand 执行命令
	ExecuteCommand(ctx context.Context, command string) (string, error)
	// ExecuteCommandWithPty 执行PTY命令
	ExecuteCommandWithPty(ctx context.Context, command string, width, height int) (io.ReadWriteCloser, error)
	// UploadFile 上传文件
	UploadFile(ctx context.Context, localPath, remotePath string) error
	// DownloadFile 下载文件
	DownloadFile(ctx context.Context, remotePath, localPath string) error
	// IsConnected 检查是否已连接
	IsConnected() bool
}

// SSHConfig SSH配置
type SSHConfig struct {
	Host           string
	Port           int
	User           string
	Password       string
	PrivateKeyPath string
	PrivateKey     []byte
	Timeout        time.Duration
}

// SSHSession SSH会话
type SSHSession struct {
	session *ssh.Session
	stdin   io.WriteCloser
	stdout  io.Reader
	stderr  io.Reader
}

// DefaultSSHClient 默认SSH客户端实现
type DefaultSSHClient struct {
	config   *SSHConfig
	client   *ssh.Client
	conn     net.Conn
	mutex    sync.RWMutex
	isClosed bool
}

// NewSSHClient 创建新的SSH客户端
func NewSSHClient(config *SSHConfig) *DefaultSSHClient {
	if config.Port == 0 {
		config.Port = 22
	}

	if config.Timeout == 0 {
		config.Timeout = 30 * time.Second
	}

	return &DefaultSSHClient{
		config: config,
	}
}

// buildSSHClientConfig 构建SSH客户端配置
func (c *DefaultSSHClient) buildSSHClientConfig() (*ssh.ClientConfig, error) {
	config := &ssh.ClientConfig{
		User:            c.config.User,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // 在生产环境中应该使用更安全的主机密钥验证
		Timeout:         c.config.Timeout,
	}

	// 如果提供了密码，则使用密码认证
	if c.config.Password != "" {
		config.Auth = append(config.Auth, ssh.Password(c.config.Password))
	}

	// 如果提供了私钥路径，则使用私钥认证
	if c.config.PrivateKeyPath != "" {
		key, err := os.ReadFile(c.config.PrivateKeyPath)
		if err != nil {
			return nil, fmt.Errorf("failed to read private key: %v", err)
		}
		signer, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("failed to parse private key: %v", err)
		}
		config.Auth = append(config.Auth, ssh.PublicKeys(signer))
	}

	// 如果提供了私钥内容，则使用私钥认证
	if len(c.config.PrivateKey) > 0 {
		signer, err := ssh.ParsePrivateKey(c.config.PrivateKey)
		if err != nil {
			return nil, fmt.Errorf("failed to parse private key: %v", err)
		}
		config.Auth = append(config.Auth, ssh.PublicKeys(signer))
	}

	// 如果没有提供认证方式，则返回错误
	if len(config.Auth) == 0 {
		return nil, fmt.Errorf("no authentication method provided")
	}

	return config, nil
}

// Connect 连接到SSH服务器
func (c *DefaultSSHClient) Connect(ctx context.Context) error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.isClosed {
		return fmt.Errorf("client is closed")
	}

	if c.client != nil {
		return nil // 已经连接
	}

	// 构建SSH客户端配置
	sshConfig, err := c.buildSSHClientConfig()
	if err != nil {
		return fmt.Errorf("failed to build SSH config: %v", err)
	}

	// 建立连接
	addr := fmt.Sprintf("%s:%d", c.config.Host, c.config.Port)
	conn, err := net.DialTimeout("tcp", addr, c.config.Timeout)
	if err != nil {
		return fmt.Errorf("failed to connect to %s: %v", addr, err)
	}
	c.conn = conn

	// 创建SSH客户端
	sshClient, chans, reqs, err := ssh.NewClientConn(conn, addr, sshConfig)
	if err != nil {
		conn.Close()
		c.conn = nil
		return fmt.Errorf("failed to create SSH connection: %v", err)
	}

	c.client = ssh.NewClient(sshClient, chans, reqs)
	return nil
}

// Disconnect 断开连接
func (c *DefaultSSHClient) Disconnect() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	c.isClosed = true

	if c.client != nil {
		err := c.client.Close()
		c.client = nil
		if c.conn != nil {
			c.conn.Close()
			c.conn = nil
		}
		return err
	}

	if c.conn != nil {
		err := c.conn.Close()
		c.conn = nil
		return err
	}

	return nil
}

// IsConnected 检查是否已连接
func (c *DefaultSSHClient) IsConnected() bool {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	return c.client != nil && !c.isClosed
}

// ExecuteCommand 执行命令
func (c *DefaultSSHClient) ExecuteCommand(ctx context.Context, command string) (string, error) {
	c.mutex.RLock()
	if c.client == nil || c.isClosed {
		c.mutex.RUnlock()
		return "", fmt.Errorf("not connected to SSH server")
	}
	client := c.client
	c.mutex.RUnlock()

	// 创建会话
	session, err := client.NewSession()
	if err != nil {
		return "", fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	// 执行命令
	output, err := session.CombinedOutput(command)
	if err != nil {
		return "", fmt.Errorf("failed to execute command: %v, output: %s", err, string(output))
	}

	return string(output), nil
}

// ExecuteCommandWithPty 执行PTY命令
func (c *DefaultSSHClient) ExecuteCommandWithPty(ctx context.Context, command string, width, height int) (io.ReadWriteCloser, error) {
	c.mutex.RLock()
	if c.client == nil || c.isClosed {
		c.mutex.RUnlock()
		return nil, fmt.Errorf("not connected to SSH server")
	}
	client := c.client
	c.mutex.RUnlock()

	// 创建会话
	session, err := client.NewSession()
	if err != nil {
		return nil, fmt.Errorf("failed to create session: %v", err)
	}

	// 请求PTY
	if err := session.RequestPty("xterm", height, width, ssh.TerminalModes{}); err != nil {
		session.Close()
		return nil, fmt.Errorf("failed to request PTY: %v", err)
	}

	// 启动命令
	if err := session.Start(command); err != nil {
		session.Close()
		return nil, fmt.Errorf("failed to start command: %v", err)
	}

	// 获取stdin和stdout
	stdin, err := session.StdinPipe()
	if err != nil {
		session.Close()
		return nil, fmt.Errorf("failed to get stdin pipe: %v", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		session.Close()
		return nil, fmt.Errorf("failed to get stdout pipe: %v", err)
	}

	// 创建读写器
	rw := &sshReadWriteCloser{
		session: session,
		stdin:   stdin,
		stdout:  stdout,
	}

	return rw, nil
}

// sshReadWriteCloser SSH读写关闭器
type sshReadWriteCloser struct {
	session *ssh.Session
	stdin   io.WriteCloser
	stdout  io.Reader
	mutex   sync.Mutex
	closed  bool
}

// Read 读取数据
func (rw *sshReadWriteCloser) Read(p []byte) (n int, err error) {
	rw.mutex.Lock()
	defer rw.mutex.Unlock()

	if rw.closed {
		return 0, io.EOF
	}

	return rw.stdout.Read(p)
}

// Write 写入数据
func (rw *sshReadWriteCloser) Write(p []byte) (n int, err error) {
	rw.mutex.Lock()
	defer rw.mutex.Unlock()

	if rw.closed {
		return 0, io.ErrClosedPipe
	}

	return rw.stdin.Write(p)
}

// Close 关闭连接
func (rw *sshReadWriteCloser) Close() error {
	rw.mutex.Lock()
	defer rw.mutex.Unlock()

	if rw.closed {
		return nil
	}

	rw.closed = true

	// 关闭stdin
	rw.stdin.Close()

	// 等待会话结束并关闭
	return rw.session.Close()
}

// UploadFile 上传文件
func (c *DefaultSSHClient) UploadFile(ctx context.Context, localPath, remotePath string) error {
	c.mutex.RLock()
	if c.client == nil || c.isClosed {
		c.mutex.RUnlock()
		return fmt.Errorf("not connected to SSH server")
	}
	client := c.client
	c.mutex.RUnlock()

	// 打开本地文件
	localFile, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("failed to open local file: %v", err)
	}
	defer localFile.Close()

	// 创建SCP会话
	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	// 获取本地文件信息
	fileInfo, err := localFile.Stat()
	if err != nil {
		return fmt.Errorf("failed to get file info: %v", err)
	}

	// 通过SCP协议上传文件
	// 这里使用一个简化的实现，实际生产环境可能需要更完整的SCP实现
	stdin, err := session.StdinPipe()
	if err != nil {
		return fmt.Errorf("failed to get stdin pipe: %v", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		return fmt.Errorf("failed to get stdout pipe: %v", err)
	}

	// 启动SCP命令
	if err := session.Start(fmt.Sprintf("scp -t %s", remotePath)); err != nil {
		return fmt.Errorf("failed to start scp command: %v", err)
	}

	// 发送文件头信息
	fmt.Fprintf(stdin, "C0644 %d %s\n", fileInfo.Size(), fileInfo.Name())

	// 发送文件内容
	_, err = io.Copy(stdin, localFile)
	if err != nil {
		return fmt.Errorf("failed to copy file content: %v", err)
	}

	// 发送结束标记
	fmt.Fprint(stdin, "\x00")
	stdin.Close()

	// 读取响应
	buf := make([]byte, 1024)
	_, err = stdout.Read(buf)
	if err != nil && err != io.EOF {
		return fmt.Errorf("failed to read response: %v", err)
	}

	// 等待会话结束
	return session.Wait()
}

// DownloadFile 下载文件
func (c *DefaultSSHClient) DownloadFile(ctx context.Context, remotePath, localPath string) error {
	c.mutex.RLock()
	if c.client == nil || c.isClosed {
		c.mutex.RUnlock()
		return fmt.Errorf("not connected to SSH server")
	}
	client := c.client
	c.mutex.RUnlock()

	// 创建本地文件
	localFile, err := os.Create(localPath)
	if err != nil {
		return fmt.Errorf("failed to create local file: %v", err)
	}
	defer localFile.Close()

	// 创建SCP会话
	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	// 通过SCP协议下载文件
	stdin, err := session.StdinPipe()
	if err != nil {
		return fmt.Errorf("failed to get stdin pipe: %v", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		return fmt.Errorf("failed to get stdout pipe: %v", err)
	}

	// 启动SCP命令
	if err := session.Start(fmt.Sprintf("scp -f %s", remotePath)); err != nil {
		return fmt.Errorf("failed to start scp command: %v", err)
	}

	// 发送确认
	stdin.Write([]byte{0})

	// 读取文件头信息
	reader := bufio.NewReader(stdout)
	_, err = reader.ReadString('\n')
	if err != nil {
		return fmt.Errorf("failed to read file header: %v", err)
	}

	// 解析文件头信息
	// 简化的实现，实际需要解析权限、大小、文件名等信息
	// 这里假设格式为: C0644 <size> <filename>

	// 发送确认
	stdin.Write([]byte{0})

	// 读取文件内容
	// 注意：这是一个简化的实现，实际SCP协议更复杂
	_, err = io.Copy(localFile, reader)
	if err != nil {
		return fmt.Errorf("failed to copy file content: %v", err)
	}

	// 发送结束确认
	stdin.Write([]byte{0})
	stdin.Close()

	// 等待会话结束
	return session.Wait()
}

// LocalSSHClient 本地SSH客户端（用于在本地执行命令）
type LocalSSHClient struct {
	workingDir string
	env        []string
}

// NewLocalSSHClient 创建新的本地SSH客户端
func NewLocalSSHClient(workingDir string, env []string) *LocalSSHClient {
	return &LocalSSHClient{
		workingDir: workingDir,
		env:        env,
	}
}

// Connect 连接（本地客户端不需要连接）
func (c *LocalSSHClient) Connect(ctx context.Context) error {
	return nil
}

// Disconnect 断开连接（本地客户端不需要断开连接）
func (c *LocalSSHClient) Disconnect() error {
	return nil
}

// ExecuteCommand 执行本地命令
func (c *LocalSSHClient) ExecuteCommand(ctx context.Context, command string) (string, error) {
	// 创建命令
	cmd := exec.CommandContext(ctx, "sh", "-c", command)
	cmd.Dir = c.workingDir
	cmd.Env = append(os.Environ(), c.env...)

	// 执行命令
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("failed to execute command: %v, output: %s", err, string(output))
	}

	return string(output), nil
}

// ExecuteCommandWithPty 执行本地PTY命令
func (c *LocalSSHClient) ExecuteCommandWithPty(ctx context.Context, command string, width, height int) (io.ReadWriteCloser, error) {
	// 创建命令
	cmd := exec.CommandContext(ctx, "sh", "-c", command)
	cmd.Dir = c.workingDir
	cmd.Env = append(os.Environ(), c.env...)

	// 获取stdin和stdout
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdin pipe: %v", err)
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdout pipe: %v", err)
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stderr pipe: %v", err)
	}

	// 启动命令
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("failed to start command: %v", err)
	}

	// 合并stdout和stderr
	reader := io.MultiReader(stdout, stderr)

	// 创建读写器
	rw := &localReadWriteCloser{
		cmd:    cmd,
		stdin:  stdin,
		stdout: reader,
	}

	return rw, nil
}

// localReadWriteCloser 本地读写关闭器
type localReadWriteCloser struct {
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	stdout io.Reader
	mutex  sync.Mutex
	closed bool
}

// Read 读取数据
func (rw *localReadWriteCloser) Read(p []byte) (n int, err error) {
	rw.mutex.Lock()
	defer rw.mutex.Unlock()

	if rw.closed {
		return 0, io.EOF
	}

	return rw.stdout.Read(p)
}

// Write 写入数据
func (rw *localReadWriteCloser) Write(p []byte) (n int, err error) {
	rw.mutex.Lock()
	defer rw.mutex.Unlock()

	if rw.closed {
		return 0, io.ErrClosedPipe
	}

	return rw.stdin.Write(p)
}

// Close 关闭连接
func (rw *localReadWriteCloser) Close() error {
	rw.mutex.Lock()
	defer rw.mutex.Unlock()

	if rw.closed {
		return nil
	}

	rw.closed = true

	// 关闭stdin
	rw.stdin.Close()

	// 终止进程
	if rw.cmd.Process != nil {
		rw.cmd.Process.Signal(syscall.SIGTERM)
		// 等待进程结束
		rw.cmd.Wait()
	}

	return nil
}

// UploadFile 上传文件（本地复制）
func (c *LocalSSHClient) UploadFile(ctx context.Context, localPath, remotePath string) error {
	// 打开源文件
	src, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("failed to open source file: %v", err)
	}
	defer src.Close()

	// 创建目标文件
	dst, err := os.Create(remotePath)
	if err != nil {
		return fmt.Errorf("failed to create destination file: %v", err)
	}
	defer dst.Close()

	// 复制文件内容
	_, err = io.Copy(dst, src)
	if err != nil {
		return fmt.Errorf("failed to copy file: %v", err)
	}

	// 同步到磁盘
	return dst.Sync()
}

// DownloadFile 下载文件（本地复制）
func (c *LocalSSHClient) DownloadFile(ctx context.Context, remotePath, localPath string) error {
	// 打开源文件
	src, err := os.Open(remotePath)
	if err != nil {
		return fmt.Errorf("failed to open source file: %v", err)
	}
	defer src.Close()

	// 创建目标文件
	dst, err := os.Create(localPath)
	if err != nil {
		return fmt.Errorf("failed to create destination file: %v", err)
	}
	defer dst.Close()

	// 复制文件内容
	_, err = io.Copy(dst, src)
	if err != nil {
		return fmt.Errorf("failed to copy file: %v", err)
	}

	// 同步到磁盘
	return dst.Sync()
}

// IsConnected 检查是否已连接（本地客户端始终连接）
func (c *LocalSSHClient) IsConnected() bool {
	return true
}
